from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals


import torch.utils.data as data
import torchvision.transforms as transforms
from PIL import Image
import PIL
import os
import os.path
import pickle
import random
import numpy as np
import pandas as pd
from miscc.config import cfg

import torch.utils.data as data
from PIL import Image
import os
import os.path
import six
import string
import sys
import torch
if sys.version_info[0] == 2:
    import cPickle as pickle
else:
    import pickle

IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG',
                  '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP']


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)


def get_imgs(img_path, imsize, bbox=None,
             transform=None, normalize=None):
    img = Image.open(img_path).convert('RGB')
    width, height = img.size
    if bbox is not None:
        r = int(np.maximum(bbox[2], bbox[3]) * 0.75)
        r = max(r, 10)
        center_x = int((2 * bbox[0] + bbox[2]) / 2)
        center_y = int((2 * bbox[1] + bbox[3]) / 2)
        y1 = np.maximum(0, center_y - r)
        y2 = np.minimum(height, center_y + r)
        x1 = np.maximum(0, center_x - r)
        x2 = np.minimum(width, center_x + r)
        img = img.crop([x1, y1, x2, y2])

    if transform is not None:
        img = transform(img)

    ret = []
    for i in range(cfg.TREE.BRANCH_NUM):
        if i < (cfg.TREE.BRANCH_NUM - 1):
            re_img = transforms.Resize(imsize[i])(img)
        else:
            re_img = img
        ret.append(normalize(re_img))

    return ret


class ImageFolder(data.Dataset):
    def __init__(self, root, split_dir='train', custom_classes=None,
                 base_size=64, transform=None, target_transform=None):
        root = os.path.join(root, split_dir)
        classes, class_to_idx = self.find_classes(root, custom_classes)
        imgs = self.make_dataset(classes, class_to_idx)
        if len(imgs) == 0:
            raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n"
                               "Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))

        self.root = root
        self.imgs = imgs
        self.classes = classes
        self.num_classes = len(classes)
        self.class_to_idx = class_to_idx

        self.transform = transform
        self.target_transform = target_transform
        self.norm = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

        self.imsize = []
        for i in range(cfg.TREE.BRANCH_NUM):
            self.imsize.append(base_size)
            base_size = base_size * 2
        print('num_classes', self.num_classes)

    def find_classes(self, dir, custom_classes):
        classes = []

        for d in os.listdir(dir):
            if os.path.isdir:
                if custom_classes is None or d in custom_classes:
                    classes.append(os.path.join(dir, d))
        print('Valid classes: ', len(classes), classes)

        classes.sort()
        class_to_idx = {classes[i]: i for i in range(len(classes))}
        return classes, class_to_idx

    def make_dataset(self, classes, class_to_idx):
        images = []
        for d in classes:
            for root, _, fnames in sorted(os.walk(d)):
                for fname in fnames:
                    if is_image_file(fname):
                        path = os.path.join(root, fname)
                        item = (path, class_to_idx[d])
                        images.append(item)
        print('The number of images: ', len(images))
        return images

    def __getitem__(self, index):
        path, target = self.imgs[index]
        imgs_list = get_imgs(path, self.imsize,
                             transform=self.transform,
                             normalize=self.norm)

        return imgs_list

    def __len__(self):
        return len(self.imgs)


class LSUNClass(data.Dataset):
    def __init__(self, db_path, base_size=64,
                 transform=None, target_transform=None):
        import lmdb
        self.db_path = db_path
        self.env = lmdb.open(db_path, max_readers=1, readonly=True, lock=False,
                             readahead=False, meminit=False)
        with self.env.begin(write=False) as txn:
            self.length = txn.stat()['entries']
            print('length: ', self.length)
        cache_file = db_path + '/cache'
        if os.path.isfile(cache_file):
            self.keys = pickle.load(open(cache_file, "rb"))
            print('Load:', cache_file, 'keys: ', len(self.keys))
        else:
            with self.env.begin(write=False) as txn:
                self.keys = [key for key, _ in txn.cursor()]
            pickle.dump(self.keys, open(cache_file, "wb"))

        self.imsize = []
        for i in range(cfg.TREE.BRANCH_NUM):
            self.imsize.append(base_size)
            base_size = base_size * 2

        self.transform = transform
        self.target_transform = target_transform
        self.norm = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    def __getitem__(self, index):
        env = self.env
        with env.begin(write=False) as txn:
            imgbuf = txn.get(self.keys[index])

        buf = six.BytesIO()
        buf.write(imgbuf)
        buf.seek(0)
        imgs = get_imgs(buf, self.imsize,
                        transform=self.transform,
                        normalize=self.norm)
        return imgs

    def __len__(self):
        return self.length

    def __repr__(self):
        return self.__class__.__name__ + ' (' + self.db_path + ')'


class TextDataset(data.Dataset):
    def __init__(self, data_dir, split='train', embedding_type='cnn-rnn',
                 base_size=64, transform=None, target_transform=None, asr_flag=False, audio_switch="0"):
        self.transform = transform
        self.norm = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        self.target_transform = target_transform

        self.imsize = []
        for i in range(cfg.TREE.BRANCH_NUM):
            self.imsize.append(base_size)
            base_size = base_size * 2

        self.data = []
        self.data_dir = data_dir
        if data_dir.find('birds') != -1:
            self.bbox = self.load_bbox()
        else:
            self.bbox = None
        split_dir = os.path.join(data_dir, split)

        self.filenames = self.load_filenames(split_dir)
        self.embeddings = self.load_embedding(split_dir, embedding_type)
        self.audio_embeddings = self.load_audio_embedding(split_dir, asr_flag, audio_switch)
        self.class_id = self.load_class_id(split_dir, len(self.filenames))
        self.captions = self.load_all_captions()

        if cfg.TRAIN.FLAG:
            self.iterator = self.prepare_training_pairs_audio  #self.prepair_training_pairs
        else:
            self.iterator = self.prepare_test_pairs_audio  #self.prepair_test_pairs

    
    def load_bbox(self):
        data_dir = self.data_dir
        bbox_path = os.path.join(data_dir, 'CUB_200_2011/bounding_boxes.txt')
        df_bounding_boxes = pd.read_csv(bbox_path,
                                        delim_whitespace=True,
                                        header=None).astype(int)
        #
        filepath = os.path.join(data_dir, 'CUB_200_2011/images.txt')
        df_filenames = \
            pd.read_csv(filepath, delim_whitespace=True, header=None)
        filenames = df_filenames[1].tolist()
        print('Total filenames: ', len(filenames), filenames[0])
        #
        filename_bbox = {img_file[:-4]: [] for img_file in filenames}
        numImgs = len(filenames)
        for i in range(0, numImgs):
            # bbox = [x-left, y-top, width, height]
            bbox = df_bounding_boxes.iloc[i][1:].tolist()

            key = filenames[i][:-4]
            filename_bbox[key] = bbox
        #
        return filename_bbox

    def load_all_captions(self):
        def load_captions(caption_name):  # self,
            cap_path = caption_name
            with open(cap_path, "rb") as f:
                captions = f.read().decode('utf8').split('\n')
            captions = [cap.replace("\ufffd\ufffd", " ")
                        for cap in captions if len(cap) > 0]
            return captions

        caption_dict = {}
        for key in self.filenames:
            caption_name = '%s/text/%s.txt' % (self.data_dir, key)
            captions = load_captions(caption_name)
            caption_dict[key] = captions
        return caption_dict

    def load_embedding(self, data_dir, embedding_type):
        if embedding_type == 'cnn-rnn':
            embedding_filename = '/char-CNN-RNN-embeddings.pickle'
        elif embedding_type == 'cnn-gru':
            embedding_filename = '/char-CNN-GRU-embeddings.pickle'
        elif embedding_type == 'skip-thought':
            embedding_filename = '/skip-thought-embeddings.pickle'

        with open(data_dir + embedding_filename, 'rb') as f:
            embeddings = pickle.load(f, encoding="bytes")
            embeddings = np.array(embeddings)
            # embedding_shape = [embeddings.shape[-1]]
            print('embeddings: ', embeddings.shape)
        return embeddings

    def load_audio_embedding(self, data_dir, asr_flag, audio_switch):
        if asr_flag:
            filename = os.path.join(data_dir, "./audio_asr_features_{}.pickle".format(audio_switch))
        else:
            filename = os.path.join(data_dir, "./audio_features_{}.pickle".format(audio_switch))
        #with open(os.path.join(data_dir, "./audio_features.pickle"), "rb") as f:
        print("audio embedding: {}".format(filename))
        with open(filename, "rb") as f:
            audio_embeddings = pickle.load(f, encoding="bytes")
            print('audio embedding:', len(audio_embeddings), audio_embeddings[0].shape)
        
        return audio_embeddings  # list

    
    def load_class_id(self, data_dir, total_num):
        if os.path.isfile(data_dir + '/class_info.pickle'):
            with open(data_dir + '/class_info.pickle', 'rb') as f:
                class_id = pickle.load(f, encoding="bytes")
        else:
            class_id = np.arange(total_num)
        return class_id

    def load_filenames(self, data_dir):
        filepath = os.path.join(data_dir, 'filenames.pickle')
        with open(filepath, 'rb') as f:
            filenames = pickle.load(f, encoding="bytes")
        print('Load filenames from: %s (%d)' % (filepath, len(filenames)))
        return filenames
    
    def prepare_training_pairs(self, index):
        key = self.filenames[index]
        if self.bbox is not None:
            bbox = self.bbox[key]
            data_dir = '%s/CUB_200_2011' % self.data_dir
        else:
            bbox = None
            data_dir = self.data_dir
        # captions = self.captions[key]
        embeddings = self.embeddings[index, :, :]
        img_name = '%s/images/%s.jpg' % (data_dir, key)
        imgs = get_imgs(img_name, self.imsize,
                        bbox, self.transform, normalize=self.norm)

        wrong_ix = random.randint(0, len(self.filenames) - 1)
        if(self.class_id[index] == self.class_id[wrong_ix]):
            wrong_ix = random.randint(0, len(self.filenames) - 1)
        wrong_key = self.filenames[wrong_ix]
        if self.bbox is not None:
            wrong_bbox = self.bbox[wrong_key]
        else:
            wrong_bbox = None
        wrong_img_name = '%s/images/%s.jpg' % \
            (data_dir, wrong_key)
        wrong_imgs = get_imgs(wrong_img_name, self.imsize,
                              wrong_bbox, self.transform, normalize=self.norm)

        embedding_ix = random.randint(0, embeddings.shape[0] - 1)
        embedding = embeddings[embedding_ix, :]
        if self.target_transform is not None:
            embedding = self.target_transform(embedding)

        return imgs, wrong_imgs, embedding, key  # captions

    
    def prepare_training_pairs_audio(self, index):
        key = self.filenames[index]
        if self.bbox is not None:
            bbox = self.bbox[key]
            data_dir = '%s/CUB_200_2011' % self.data_dir
        else:
            bbox = None
            data_dir = self.data_dir
        # captions = self.captions[key]
        embeddings = self.audio_embeddings[index]  # list

        img_name = '%s/images/%s.jpg' % (data_dir, key)
        imgs = get_imgs(img_name, self.imsize,
                        bbox, self.transform, normalize=self.norm)

        wrong_ix = random.randint(0, len(self.filenames) - 1)
        if(self.class_id[index] == self.class_id[wrong_ix]):
            wrong_ix = random.randint(0, len(self.filenames) - 1)
        wrong_key = self.filenames[wrong_ix]
        if self.bbox is not None:
            wrong_bbox = self.bbox[wrong_key]
        else:
            wrong_bbox = None
        wrong_img_name = '%s/images/%s.jpg' % \
            (data_dir, wrong_key)
        wrong_imgs = get_imgs(wrong_img_name, self.imsize,
                              wrong_bbox, self.transform, normalize=self.norm)

        embedding_ix = random.randint(0, len(embeddings) - 1)
        embedding = embeddings[embedding_ix]
        if self.target_transform is not None:
            embedding = self.target_transform(embedding)

        return imgs, wrong_imgs, embedding, key, self.class_id[index]  # captions

    def prepare_test_pairs(self, index):
        key = self.filenames[index]
        if self.bbox is not None:
            bbox = self.bbox[key]
            data_dir = '%s/CUB_200_2011' % self.data_dir
        else:
            bbox = None
            data_dir = self.data_dir
        # captions = self.captions[key]
        embeddings = self.embeddings[index, :, :]
        img_name = '%s/images/%s.jpg' % (data_dir, key)
        imgs = get_imgs(img_name, self.imsize,
                        bbox, self.transform, normalize=self.norm)

        if self.target_transform is not None:
            embeddings = self.target_transform(embeddings)

        return imgs, embeddings, key  # captions

    
    def prepare_test_pairs_audio(self, index):
        key = self.filenames[index]
        if self.bbox is not None:
            bbox = self.bbox[key]
            data_dir = '%s/CUB_200_2011' % self.data_dir
        else:
            bbox = None
            data_dir = self.data_dir
        # captions = self.captions[key]
        embeddings = self.audio_embeddings[index]
        img_name = '%s/images/%s.jpg' % (data_dir, key)
        imgs = get_imgs(img_name, self.imsize,
                        bbox, self.transform, normalize=self.norm)
        if isinstance(embeddings, list):
            if len(embeddings)<10:
                embeddings += [embeddings[-1]]*(10-len(embeddings))
        
        if self.target_transform is not None:
            embeddings = self.target_transform(embeddings)
        embeddings = torch.Tensor(embeddings)
        return imgs, embeddings, key  # captions

    def __getitem__(self, index):
        return self.iterator(index)

    def __len__(self):
        return len(self.filenames)


import json

class BaseDataset(data.Dataset):
    def __init__(self, data_root, train=True, base_size=64, transform=None, target_transform=None, feature_switch='image'):
        split = "train" if train else "test"
        self.data_root = data_root
        self.json_data_all = json.load(open(os.path.join(self.data_root, "{}.json".format(split))))
        self.image_folder = self.json_data_all["image_base_path"]
        self.audio_folder = self.json_data_all['audio_base_path']
        self.json_data = self.json_data_all['data']
        embedding_path = os.path.join(data_root, split, "audio_features_{}.pickle".format(feature_switch))
        print("load features from: {}".format(embedding_path))
        self.embedding = pickle.load(open(embedding_path, "rb"))
        
        self.imsize = []
        for i in range(cfg.TREE.BRANCH_NUM):
            self.imsize.append(base_size)
            base_size = base_size * 2

        self.transform = transform
        self.target_transform = target_transform
        self.norm = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        self.iterator = self.prepare_train_pairs if train else self.prepare_test_pairs

    def __len__(self):
        return len(self.json_data)
    
    def _get_img(self, item):
        return item['img']

    def _get_class(self, item):
        return int(item['class'])

    def find_wrong_image(self, base_class_label):
        while True:
            json_data, _ = self.get_rand(self.json_data)
            image_path, class_label = self._get_img(json_data), self._get_class(json_data)
            if class_label != base_class_label:
                break
        return image_path
    
    @staticmethod
    def get_rand(feature):
        # feature: (10, 1024)
        rand_idx = random.randint(0, len(feature)-1)
        return feature[rand_idx], rand_idx

    def prepare_train_pairs(self, index):
        json_data = self.json_data[index]
        image_path, class_label = self._get_img(json_data), self._get_class(json_data)
        embedding, _ = self.get_rand(self.embedding[index])
        wrong_image_path = self.find_wrong_image(class_label)
        real_image = get_imgs(
            os.path.join(self.image_folder, image_path), self.imsize, transform=self.transform, normalize=self.norm
        )
        wrong_image = get_imgs(
            os.path.join(self.image_folder, wrong_image_path), self.imsize, transform=self.transform, normalize=self.norm
        )
        return real_image, wrong_image, embedding, image_path, class_label

    def prepare_test_pairs(self, index):
        json_data = self.json_data[index]
        image_path, class_label = self._get_img(json_data), self._get_class(json_data)
        # embedding, _ = self.get_rand(self.embedding[index])
        embedding = self.embedding[index]

        real_image = get_imgs(
            os.path.join(self.image_folder, image_path), self.imsize, transform=self.transform, normalize=self.norm
        )
        return real_image, embedding, image_path

    def __getitem__(self, index):
        return self.iterator(index)


   
class BirdsDataset(BaseDataset):
    def __init__(self, data_root, train=True, base_size=64, transform=None, target_transform=None, feature_switch='image'):
        super(BirdsDataset, self).__init__(data_root, train, base_size, transform, target_transform, feature_switch)
        self.bbox = self.load_bbox()

    def load_bbox(self):
        data_dir = self.data_root
        bbox_path = os.path.join(data_dir, 'CUB_200_2011/bounding_boxes.txt')
        df_bounding_boxes = pd.read_csv(bbox_path,
                                        delim_whitespace=True,
                                        header=None).astype(int)
        #
        filepath = os.path.join(data_dir, 'CUB_200_2011/images.txt')
        df_filenames = \
            pd.read_csv(filepath, delim_whitespace=True, header=None)
        filenames = df_filenames[1].tolist()
        print('Total filenames: ', len(filenames), filenames[0])
        #
        filename_bbox = {img_file[:-4]: [] for img_file in filenames}
        numImgs = len(filenames)
        for i in range(0, numImgs):
            # bbox = [x-left, y-top, width, height]
            bbox = df_bounding_boxes.iloc[i][1:].tolist()

            key = filenames[i][:-4]
            filename_bbox[key] = bbox
        #
        return filename_bbox

    def _get_img(self, item):
        return item['image']

    def _get_class(self, item):
        return int(item['class'].split('.')[0])
    
    def _get_bbox(self, image_path):
        return self.bbox[image_path[:-4]]

    def prepare_train_pairs(self, index):
        json_data = self.json_data[index]
        image_path, class_label = self._get_img(json_data), self._get_class(json_data)
        embedding, _ = self.get_rand(self.embedding[index])
        wrong_image_path = self.find_wrong_image(class_label)
        real_bbox = self._get_bbox(image_path)
        wrong_bbox = self._get_bbox(wrong_image_path)
        
        real_image = get_imgs(
            os.path.join(self.image_folder, image_path), self.imsize, bbox=real_bbox, transform=self.transform, normalize=self.norm
        )
        wrong_image = get_imgs(
            os.path.join(self.image_folder, wrong_image_path), self.imsize, bbox=wrong_bbox, transform=self.transform, normalize=self.norm
        )
        return real_image, wrong_image, embedding, image_path, class_label

    def prepare_test_pairs(self, index):
        json_data = self.json_data[index]
        image_path, class_label = self._get_img(json_data), self._get_class(json_data)
        # embedding, _ = self.get_rand(self.embedding[index])
        embedding = self.embedding[index]
        real_image = get_imgs(
            os.path.join(self.image_folder, image_path), self.imsize, transform=self.transform, normalize=self.norm
        )
        return real_image, embedding, image_path

    def __getitem__(self, index):
        return self.iterator(index)


class PlacesSubSet(BaseDataset):
    class_label = {
        'bedroom':0, 'dinette':1, 'dining_room':2, 'home_office':3, 'hotel_room':4,
        'kitchenette':5, 'living_room':6
    }
    def __init__(self, data_root, train=True, base_size=64, transform=None, target_transform=None, feature_switch="cnn_googlenet"):
        super(PlacesSubSet, self).__init__(data_root, train, base_size, transform, target_transform, feature_switch)
    
    def _get_img(self, item):
        return item['image'][2:]

    def _get_class(self, item):
        return self.class_label[item['image'].split("/")[1]]

    def prepare_training_pairs(self, index):
        json_data = self.json_data[index]
        image_path, label = self._get_img(json_data), self._get_class(json_data)
        embedding = self.image_embedding[index]
        wrong_image_path = self.find_wrong_image(label)
        #print(self.data_root, image_path, wrong_image_path)
        real_image = get_imgs(
            os.path.join(self.json_data['image_base_path'], image_path), self.imsize, transform=self.transform, normalize=self.norm
            )
        wrong_image = get_imgs(
            os.path.join(self.json_data['image_base_path'], wrong_image_path), self.imsize, transform=self.transform, normalize=self.norm
            ) 
        return real_image, wrong_image, embedding, image_path, label

    def prepare_test_pairs(self, index):
        json_data = self.json_data[index]
        image_path, label = self._get_img(json_data), self._get_class(json_data)
        embedding = np.expand_dims(self.image_embedding[index], axis=0)
        #print(self.data_root, image_path, wrong_image_path)
        real_image = get_imgs(
            os.path.join(self.json_data['image_base_path'], image_path), self.imsize, transform=self.transform, normalize=self.norm
            )
        
        return real_image, embedding, image_path

   
class FlowersDataset(data.Dataset):
    def __init__(self, data_root, train=True, base_size=64, transform=None, target_transform=None, feature_switch='image'):
        super(FlowersDataset, self).__init__(data_root, train, base_size, transform, target_transform, feature_switch)

    def _get_img(self, item):
        return item['img']

    def _get_class(self, item):
        return int(item['class'])

    def prepare_train_pairs(self, index):
        json_data = self.json_data[index]
        image_path, class_label = self._get_img(json_data), self._get_class(json_data)
        embedding, _ = self.get_rand(self.embedding[index])
        wrong_image_path = self.find_wrong_image(class_label)
        real_image = get_imgs(
            os.path.join(self.image_folder, image_path), self.imsize, transform=self.transform, normalize=self.norm
        )
        wrong_image = get_imgs(
            os.path.join(self.image_folder, wrong_image_path), self.imsize, transform=self.transform, normalize=self.norm
        )
        return real_image, wrong_image, embedding, image_path, class_label

    def prepare_test_pairs(self, index):
        json_data = self.json_data[index]
        image_path, class_label = self._get_img(json_data), self._get_class(json_data)
        # embedding, _ = self.get_rand(self.embedding[index])
        embedding = self.embedding[index]
        real_image = get_imgs(
            os.path.join(self.image_folder, image_path), self.imsize, transform=self.transform, normalize=self.norm
        )
        return real_image, embedding, image_path

    def __getitem__(self, index):
        return self.iterator(index)

